from depthaware.options.train_options import TrainOptions
from depthaware.data.datasets_SSD_nyu import nyuSSDDataset
from depthaware.models.models import create_model
from depthaware.utils import util
from depthaware.utils.utils import *
from depthaware.utils.visualizer import Visualizer
from depthaware.models.SSD_Deeplab import SSD, MultiBoxLoss
from tqdm import tqdm
import os
import numpy as np
import time
import torch

DEBUG = False
if DEBUG:
    np.random.seed(1)
    torch.manual_seed(0)
else:
    np.random.seed(int(time.time()))

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
grad_clip = None
print_freq = 200  # print training status every __ batches
workers = 4  # number of workers for loading data in the DataLoader

def train(train_loader, model, criterion, optimizer, epoch):
    """
    One epoch's training.

    :param train_loader: DataLoader for training data
    :param model: model
    :param criterion: MultiBox loss
    :param optimizer: optimizer
    :param epoch: epoch number
    """
    model.train()  # training mode enables dropout

    batch_time = AverageMeter()  # forward prop. + back prop. time
    data_time = AverageMeter()  # data loading time
    losses = AverageMeter()  # loss

    start = time.time()

    # Batches
    for i, (images, boxes, clas, depths, _) in enumerate(train_loader):
        data_time.update(time.time() - start)

        # Move to default device
        images = images.to(device)  # (batch_size (N), 3, 300, 300)
        boxes = [b.to(device) for b in boxes]
        clas = [c.to(device) for c in clas]
        depths = depths.to(device)  # (batch_size (N), 1, 300, 300)


        # Forward prop.
        predicted_locs, predicted_scores = model(images, depths)  # (N, 8732, 4), (N, 8732, n_classes)

        # Loss
        loss = criterion(predicted_locs, predicted_scores, boxes, clas)  # scalar

        # Backward prop.
        optimizer.zero_grad()
        loss.backward()

        # Clip gradients, if necessary
        if grad_clip is not None:
            clip_gradient(optimizer, grad_clip)

        # Update model
        optimizer.step()

        losses.update(loss.item(), images.size(0))
        batch_time.update(time.time() - start)

        start = time.time()

        # Print status
        if i % print_freq == 0:
            print('Epoch: [{0}][{1}/{2}]\t'
                  'Batch Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
                  'Data Time {data_time.val:.3f} ({data_time.avg:.3f})\t'
                  'Loss {loss.val:.4f} ({loss.avg:.4f})\t'.format(epoch, i, len(train_loader),
                                                                  batch_time=batch_time,
                                                                  data_time=data_time, loss=losses))
    del predicted_locs, predicted_scores, images, boxes, clas  # free some memory since their histories may be stored
    # print('Epoch: [{0}][{1}/{2}]\t'
    #       'Batch Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
    #       'Data Time {data_time.val:.3f} ({data_time.avg:.3f})\t'
    #       'Loss {loss.val:.4f} ({loss.avg:.4f})\t'.format(epoch, i, len(train_loader),
    #                                                       batch_time=batch_time,
    #                                                       data_time=data_time, loss=losses))


if __name__ == '__main__':
    torch.multiprocessing.freeze_support()

    opt = TrainOptions().parse()
    n_classes = opt.label_nc
    lr = opt.lr
    momentum = opt.momentum
    weight_decay = opt.wd
    depthconv = opt.depthconv
    batchsize = opt.batchSize
    checkpoint_path = os.path.join(opt.checkpoints_dir, opt.name, 'checkpoint.pt')
    if opt.continue_train:
        checkpoint = torch.load(checkpoint_path)
        start_epoch = checkpoint['epoch'] + 1
        print('\nLoaded checkpoint from epoch %d.\n' % start_epoch)
        model = checkpoint['model']
        optimizer = checkpoint['optimizer']
    else:
        start_epoch = 0
        model = SSD(n_classes=n_classes, depthconv=depthconv)
        # Initialize the optimizer, with twice the default learning rate for biases, as in the original Caffe repo
        biases = list()
        not_biases = list()
        for param_name, param in model.named_parameters():
            if param.requires_grad:
                if param_name.endswith('.bias'):
                    biases.append(param)
                else:
                    not_biases.append(param)
        optimizer = torch.optim.SGD(params=[{'params': biases, 'lr': 2 * lr}, {'params': not_biases}],
                                    lr=lr, momentum=momentum, weight_decay=weight_decay)

    # Move to default device
    model = model.to(device)
    criterion = MultiBoxLoss(priors_cxcy=model.priors_cxcy).to(device)

    # data_loader = CreateDataLoader(opt)
    dataset = nyuSSDDataset(opt, split='train')
    # dataset, dataset_val = data_loader.load_data()
    dataset_size = len(dataset)
    train_loader = torch.utils.data.DataLoader(dataset, batch_size=batchsize, shuffle=True,
                                               collate_fn=dataset.collate_fn, num_workers=workers,
                                               pin_memory=True, drop_last=True)  # note that we're passing the collate function here
    print('#training images = %d' % dataset_size)

    epochs = opt.nepochs
    decay_lr_at = [epochs*2/3, epochs*5/6]
    for epoch in range(start_epoch, epochs):

        # Decay learning rate at particular epochs
        if epoch in decay_lr_at:
            adjust_learning_rate(optimizer, decay_lr_at)

            # One epoch's training
        train(train_loader=train_loader,
              model=model,
              criterion=criterion,
              optimizer=optimizer,
              epoch=epoch)

        save_checkpoint(epoch, model, optimizer)


